import os
import cv2
import numpy as np


def upsample_flow(flow, scale_factor, interpolation=cv2.INTER_CUBIC):
    h, w = flow.shape[:2]
    new_h, new_w = int(h * scale_factor), int(w * scale_factor)

    # 上采样光流的每个分量
    flow_up_x = cv2.resize(flow[:, :, 0], (new_w, new_h), interpolation=interpolation) * scale_factor
    flow_up_y = cv2.resize(flow[:, :, 1], (new_w, new_h), interpolation=interpolation) * scale_factor
    flow_up = np.stack((flow_up_x, flow_up_y), axis=2)

    return flow_up


def save_flow_part(flow_part, row, col, output_dir):
    filename = f"tr{row + 1}-tc{col + 1}.npy"
    filepath = os.path.join(output_dir, filename)
    np.save(filepath, flow_part)
    print(f"Saved {filename} to {filepath}")


def split_and_upsample_flow(flow, small_tile_size, large_tile_size, output_dir, scale_factor):
    h, w = flow.shape[:2]
    os.makedirs(output_dir, exist_ok=True)

    for row in range(0, h, small_tile_size):
        for col in range(0, w, small_tile_size):
            # 提取当前块
            flow_part = flow[row:row + small_tile_size, col:col + small_tile_size]

            # 上采样当前块
            flow_part_up = upsample_flow(flow_part, scale_factor)

            # 保存上采样后的块
            save_flow_part(flow_part_up, row // small_tile_size, col // small_tile_size, output_dir)


# 示例用法
input_flow_path = 'E:\\wafer52\\256nm_flow_combine\\large_flow_repair3.npy'  # 替换为实际输入光流场文件路径
output_dir = 'E:\\wafer52\\coarse_32nm_flow3\\'  # 替换为实际输出目录

# 加载256nm的光流场
flow_256nm = np.load(input_flow_path)

# 以128x128为单位分割，并上采样至1024x1024
split_and_upsample_flow(flow_256nm, small_tile_size=128, large_tile_size=1024, output_dir=output_dir,
                        scale_factor=1024 / 128)
